﻿#include <iostream>
#include <fstream>
#include <thread>

#include <zmq.hpp>

#include "../../inc/helper.hpp"
#include "nlohmann/json.hpp"

using nlohmann::json;

const static int g_clients_number = 2;
const static int g_worker_number = 3;
const static char* g_worker_ready_msg = "READY";

struct DcConfig
{
	int m_state_port;
	int m_cloud_frontend_port;
};

std::string read_file(std::string file_name)
{
	std::ifstream file;
	file.open(file_name, std::ios_base::in | std::ios_base::out);
	if (!file)
	{
		info("failed to open file {}", file_name);
		return "";
	}

	file.seekg(0, std::ios_base::end);
	size_t file_size = static_cast<size_t>(file.tellg());
	file.seekg(0, std::ios_base::beg);

	std::string content;
	content.resize(file_size);
	if (file_size <= 0)
	{
		info("empty file: {}", file_name);
		return "";
	}

	file.read(&content[0], file_size);
	return content;
}

DcConfig parse_config(std::string content)
{
	auto json_obj = json::parse(content);
	DcConfig cfg;
	cfg.m_state_port = json_obj["state_port"].get<int>();
	cfg.m_cloud_frontend_port = json_obj["m_cloud_frontend_port"].get<int>();
	return cfg;
}

DcConfig get_dc_config(std::string dc_name)
{
	std::string cfg_file = dc_name + ".json";
	std::string content = read_file(cfg_file);
	return parse_config(content);
}

void client_task(zmq::context_t& ctx)
{
	zmq::socket_t client(ctx, zmq::socket_type::req);
	client.connect("inproc://localfe");

	std::random_device rd;
	std::uniform_int_distribution<int> dist(0, 4);
	std::uniform_int_distribution<int> burst_gen(0, 14);
	std::uniform_int_distribution<int> task_id_gen(0, 0x10000);
	while (true)
	{
		std::this_thread::sleep_for(std::chrono::seconds(dist(rd)));
		int burst = burst_gen(rd);
		while (burst--)
		{
			std::string task_id = fmt::format("{:04X}", task_id_gen(rd));
			send_string(client, task_id);

			zmq::pollitem_t item = { client, 0, ZMQ_POLLIN, 0 };
			zmq::poll(&item, 1, 10 * 1000);
			if (item.revents & ZMQ_POLLIN)
			{
				auto msgs = zmsg_recv(client);
				info(to_string(msgs.back()));
			}
			else
			{
				info("CLIENT EXIT - lost task {}", task_id);
			}
		}
	}
}

void worker_task(zmq::context_t& ctx)
{
	zmq::socket_t worker(ctx, zmq::socket_type::req);
	worker.connect("inproc://localbe");

	worker.send(to_message(g_worker_ready_msg));

	std::random_device rd;
	std::uniform_int_distribution<int> dist(0, 1);
	while (true)
	{
		auto msgs = zmsg_recv(worker);

		std::this_thread::sleep_for(std::chrono::seconds(dist(rd)));
		
		zmsg_send(worker, msgs);
	}
}

int main(int argc, char* argv[])
{
	if (argc < 2)
	{
		info("syntax: peering3 me {you}...");
		return 1;
	}
	std::string self = argv[1];
	info("preparing broker at {}...", self);

	DcConfig my_cfg = get_dc_config(self);

	zmq::context_t ctx;
	zmq::socket_t cloudfe(ctx, zmq::socket_type::router);
	cloudfe.setsockopt(ZMQ_IDENTITY, self.data(), self.size());
	cloudfe.bind(fmt::format("tcp://*:{}", my_cfg.m_cloud_frontend_port));

	zmq::socket_t localfe(ctx, zmq::socket_type::router);
	localfe.bind("inproc://localfe");
	zmq::socket_t localbe(ctx, zmq::socket_type::router);
	localbe.bind("inproc://localbe");

	zmq::socket_t cloudbe(ctx, zmq::socket_type::router);
	cloudbe.setsockopt(ZMQ_IDENTITY, self.data(), self.size());
	for (int i = 2; i < argc; ++i)
	{
		std::string peer = argv[i];
		DcConfig peer_cfg = get_dc_config(peer);
		info("connecting to cloud frontend at {}", peer);
		cloudbe.connect(fmt::format("tcp://localhost:{}", peer_cfg.m_cloud_frontend_port));
	}

	zmq::socket_t statebe(ctx, zmq::socket_type::pub);
	statebe.bind(fmt::format("tcp://*:{}", my_cfg.m_state_port));

	zmq::socket_t statefe(ctx, zmq::socket_type::sub);
	statefe.setsockopt(ZMQ_SUBSCRIBE, "", 0);
	for (int i = 2; i < argc; ++i)
	{
		DcConfig dc_config = get_dc_config(argv[i]);
		info("connecting to state backend at {}-{}", argv[i], dc_config.m_state_port);
		statefe.connect(fmt::format("tcp://localhost:{}", dc_config.m_state_port));
	}

	info("Press any key when all brokers are started: ");
	pause();

	for (int i = 0; i < g_worker_number; ++i)
	{
		std::thread t(worker_task, std::ref(ctx));
		t.detach();
	}

	for (int i = 0; i < g_clients_number; ++i)
	{
		std::thread t(client_task, std::ref(ctx));
		t.detach();
	}

	int cloud_capacity = 0;
	std::list<zmq::message_t> worker_list;
	std::vector<zmq::pollitem_t> primary = {
		{ localbe, 0, ZMQ_POLLIN, 0 },
		{ cloudbe, 0, ZMQ_POLLIN, 0 },
		{ statefe, 0, ZMQ_POLLIN, 0 }
	};

	while (true)
	{
		// Track if capacity changes during this iteration
		int previous = worker_list.size();

		zmq::poll(primary, worker_list.empty() ? -1 : 100);
		std::deque<zmq::message_t> msgs;
		if (primary[0].revents & ZMQ_POLLIN)
		{
			msgs = zmsg_recv(localbe);
			worker_list.push_back(std::move(msgs.front()));
			msgs.pop_front();
			msgs.pop_front();
			if (to_string(msgs.back()) == g_worker_ready_msg)
			{
				msgs.clear();
			}
		}
		else if (primary[1].revents & ZMQ_POLLIN)
		{
			msgs = zmsg_recv(cloudbe);
			info("reply from cloud");
			msgs.pop_front();
			msgs.pop_front();
		}

		if (!msgs.empty())
		{
			for (int i = 2; i < argc; ++i)
			{
				if (to_string(msgs.front()) == argv[i])
				{
					zmsg_send(cloudfe, msgs);
					msgs.clear();
					break;
				}
			}

			if (!msgs.empty())
			{
				zmsg_send(localfe, msgs);
			}
		}

		if (primary[2].revents & ZMQ_POLLIN)
		{
			msgs = zmsg_recv(statefe);
			std::string peer_name = to_string(msgs[0]);
			std::string available = to_string(msgs[1]);
			cloud_capacity = std::atoi(available.c_str());
		}

		while (worker_list.size()+cloud_capacity > 0)
		{
			std::vector<zmq::pollitem_t> secondary = {
				{ localfe, 0, ZMQ_POLLIN, 0 },
				{ cloudfe, 0, ZMQ_POLLIN, 0 }
			};

			if (worker_list.size() > 0)
			{
				zmq::poll(secondary, 0);
			}
			else
			{
				zmq::poll(secondary.data(), 1, 0);
			}
			
			if (secondary[1].revents & ZMQ_POLLIN)
			{
				msgs = zmsg_recv(cloudfe);
			}
			else if (secondary[0].revents & ZMQ_POLLIN)
			{
				msgs = zmsg_recv(localfe);
			}
			else
			{
				break;
			}

			if (worker_list.size() > 0)
			{
				msgs.push_front(zmq::message_t());
				msgs.push_front(std::move(worker_list.front()));
				worker_list.pop_front();
				zmsg_send(localbe, msgs);
			}
			else
			{
				int random_peer = rand() % (argc - 2) + 2;
				msgs.push_front(zmq::message_t());
				msgs.push_front(to_message(argv[random_peer]));
				zmsg_send(cloudbe, msgs);
			}
		}

		if (worker_list.size() != previous)
		{
			std::deque<zmq::message_t> msgs;
			msgs.emplace_back(to_message(self));
			msgs.emplace_back(to_message(fmt::format("{}", rand() % 10)));
			zmsg_send(statebe, msgs);
		}
	}
}